import cv2
import os
import logging
import argparse

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s: %(message)s')
parser = argparse.ArgumentParser(description='Video to Pictures')
parser.add_argument('--mode', type=int, default=0, help='mode of video.(0:Offline Video; 1:RTSP; 2:Camera)')
parser.add_argument('--path', type=str, help='offline video path or rtsp video stream link')
parser.add_argument('--cnt', type=int, default=0, help='image start number')
parser.add_argument('--span', type=int, default=1, help='time span, default to 1 second')
parser.add_argument('--limit', type=int, default=1000, help='upper limit of saved images')
parser.add_argument('--show', type=int, default=0, help='whether to display')

args = parser.parse_args()

def read_rtsp():
    if args.mode == 0 or args.mode == 1:
        if args.path is None:
            logging.info("Please enter Path.")
            return
        cap = cv2.VideoCapture(args.path)  # 读取视频
    elif args.mode == 2:
        cap = cv2.VideoCapture(0)
    else:
        logging.info("Mode can only be 0, 1, 2. Please confirm if the input parameters are correct.")
        return

    v_size = str(int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))) + " X " + str(int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    v_fps = str(int(cap.get(cv2.CAP_PROP_FPS)))
    # fourcc = cv2.VideoWriter_fourcc(*'XVID')    # cv2.VideoWriter_fourcc('X', 'V', 'I', 'D')表示 MPEG-4 编码类型
    # out = cv2.VideoWriter('./rtsp_video.avi', fourcc, fps, size)  #视频流存储本地

    logging.info("Size: " + v_size + ", FPS: " + v_fps)

    count = args.cnt  # 图片标号
    frame_interval = int(args.span * v_fps)  # 帧数
    frame_interval_count = 0

    #检查是否打开正确
    if cap.isOpened():
        open, frame = cap.read()
    else:
        open = False

    while open:
        # 读取每一帧
        ret, frame = cap.read()
        if frame is None:
            break
        if ret == True:
            if args.show == 1:
                cv2.imshow('frame', frame)
                cv2.waitKey(1)
            if frame_interval_count % frame_interval == 0:
                save_image(count, frame)
                logging.info("num：" + str(count) + ", frame: " + str(frame_interval_count))
                count += 1
                if count > args.limit:
                    break
            frame_interval_count += 1
            # out.write(frame)

    cap.release()
    # out.release()
    cv2.destroyAllWindows()

def save_image(num, image):

    dir_path = 'raw_pictures'
    if not os.path.exists(dir_path):
        os.mkdir(dir_path)
    image_path = dir_path + '/{}.jpg'.format(str(num))
    print(image_path)
    cv2.imwrite(image_path, image)

if __name__ == '__main__':
    read_rtsp()
